package com.zt.chatglm.session.defaults;

import com.zt.chatglm.commons.Constants;
import com.zt.chatglm.domain.chat.ChatCompletionRequest;
import com.zt.chatglm.domain.chat.ChatCompletionSyncResponse;
import com.zt.chatglm.domain.image.ImageCompletionRequest;
import com.zt.chatglm.domain.image.ImageCompletionResponse;
import com.zt.chatglm.executor.Executor;
import com.zt.chatglm.session.Configuration;
import com.zt.chatglm.session.OpenAiSession;
import lombok.extern.slf4j.Slf4j;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;

import java.util.Map;
import java.util.concurrent.CompletableFuture;

/**
 * @author: zt
 * @Description: 会话服务实现
 * @create: 2024-05-20 23:43
 */

@Slf4j
public class DefaultOpenAiSession implements OpenAiSession {

    private final Configuration configuration;
    private final Map<Constants.Model, Executor> executorGroup;

    public DefaultOpenAiSession(Configuration configuration, Map<Constants.Model, Executor> executorGroup) {
        this.configuration = configuration;
        this.executorGroup = executorGroup;
    }

    @Override
    public EventSource completions(ChatCompletionRequest completionRequest, EventSourceListener eventSourceListener) throws Exception {
        Executor executor = executorGroup.get(completionRequest.getModel());
        if (null == executor){
            throw new RuntimeException(completionRequest.getModel() + " 模型执行器尚未实现！");
        }
        return executor.completions(completionRequest, eventSourceListener);
    }

    @Override
    public CompletableFuture<String> completions(ChatCompletionRequest completionRequest) throws Exception {
        Executor executor = executorGroup.get(completionRequest.getModel());
        if (null == executor) {
            throw new RuntimeException(completionRequest.getModel() + " 模型执行器尚未实现！");
        }
        return executor.completions(completionRequest);
    }

    @Override
    public ChatCompletionSyncResponse completionsSync(ChatCompletionRequest completionRequest) throws Exception {
        Executor executor = executorGroup.get(completionRequest.getModel());
        if (null == executor) {
            throw new RuntimeException(completionRequest.getModel() + " 模型执行器尚未实现！");
        }
        return executor.completionsSync(completionRequest);
    }

    @Override
    public ImageCompletionResponse genImages(ImageCompletionRequest request) throws Exception {
        Executor executor = executorGroup.get(request.getModelEnum());
        if (null == executor) {
            throw new RuntimeException(request.getModelEnum() + " 模型执行器尚未实现！");
        }
        return executor.genImages(request);
    }

    @Override
    public Configuration getConfiguration() {
        return configuration;
    }
}
